from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import mysql.connector
import time

# 数据库连接设置
db_config = {
    'user': 'root',
    'password': 'Wlj98192188?',
    'host': 'localhost',
    'database': 'stock_data'
}

# 初始化数据库连接
def init_db():
    conn = mysql.connector.connect(**db_config)
    cursor = conn.cursor()
    return conn, cursor

# 插入数据到数据库
def insert_data(cursor, conn, data):
    sql = """
    INSERT INTO stocks_2 (
        bStockNo, bStockName, bLatestPrice, bChangePercent, bChangeAmount, 
        bVolume, bTurnover, bAmplitude, bHigh, bLow, bOpen, bPreviousClose
    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    """
    cursor.execute(sql, data)
    conn.commit()

# 配置 Selenium
options = webdriver.ChromeOptions()
options.add_argument('--headless')  # 无头模式
driver = webdriver.Chrome(options=options)
wait = WebDriverWait(driver, 10)

# 处理浮点数转换的函数
def clean_decimal(value):
    if not value:
        return None
    value = value.strip().replace(',', '')
    try:
        if '%' in value:
            return float(value.replace('%', '')) / 100
        elif '亿' in value:
            return float(value.replace('亿', '')) * 1e8
        elif '万' in value:
            return float(value.replace('万', '')) * 1e4
        return float(value)
    except ValueError:
        return None

# 爬取数据的函数
def scrape_data():
    stock_data = []

    # 等待表格加载并找到所有行
    wait.until(EC.presence_of_element_located((By.XPATH, '//table//tr')))
    rows = driver.find_elements(By.XPATH, '//table//tr')

    for row in rows:
        cols = row.find_elements(By.TAG_NAME, 'td')
        if len(cols) >= 14:
            bStockNo = cols[1].text.strip()
            bStockName = cols[2].text.strip()
            bLatestPrice = clean_decimal(cols[4].text.strip())
            bChangePercent = clean_decimal(cols[5].text.strip())
            bChangeAmount = clean_decimal(cols[6].text.strip())
            bVolume = clean_decimal(cols[7].text.strip())
            bTurnover = clean_decimal(cols[8].text.strip())
            bAmplitude = clean_decimal(cols[9].text.strip())
            bHigh = clean_decimal(cols[10].text.strip())
            bLow = clean_decimal(cols[11].text.strip())
            bOpen = clean_decimal(cols[12].text.strip())
            bPreviousClose = clean_decimal(cols[13].text.strip())

            stock_info = (
                bStockNo, bStockName, bLatestPrice, bChangePercent, bChangeAmount,
                bVolume, bTurnover, bAmplitude, bHigh, bLow, bOpen, bPreviousClose
            )

            # 打印每条数据
            print(stock_info)

            # 添加数据到列表中
            stock_data.append(stock_info)

    return stock_data

# 模拟点击切换板块
def switch_section(section_id):
    section_tab = driver.find_element(By.ID, section_id)
    section_tab.click()
    time.sleep(3)  # 等待页面加载

# 主函数
def main():
    url = "http://quote.eastmoney.com/center/gridlist.html#hs_a_board"

    # 连接数据库
    conn, cursor = init_db()

    try:
        # 打开网页
        driver.get(url)

        # 爬取沪深A股板块数据
        print("开始爬取沪深A股板块数据...")
        hs_data = scrape_data()
        for stock in hs_data:
            insert_data(cursor, conn, stock)
        print("沪深A股数据爬取并存储成功！")

        # 切换到上证A股板块并爬取数据
        print("切换到上证A股板块...")
        switch_section("nav_sh_a_board")
        sz_data = scrape_data()
        for stock in sz_data:
            insert_data(cursor, conn, stock)
        print("上证A股数据爬取并存储成功！")

        # 切换到深证A股板块并爬取数据
        print("切换到深证A股板块...")
        switch_section("nav_sz_a_board")
        zx_data = scrape_data()
        for stock in zx_data:
            insert_data(cursor, conn, stock)
        print("深证A股数据爬取并存储成功！")

    except Exception as e:
        print("爬取失败：", e)

    finally:
        cursor.close()
        conn.close()
        driver.quit()

if __name__ == "__main__":
    main()
